#include <iostream>
#include <fstream>
#include <vector>
#include <cstdint>
#include <string>
#include <nlohmann/json.hpp>
#include "pointcloud_tools.h"
#include "point_preprocess.h"
#include "point_common.h"

using json = nlohmann::json;
using namespace pointcloud_tools;
using namespace PointPreprocess;

namespace model_builder
{
    enum class LayerType
    {
        SUBMCONV,
        FC
    };

    struct LayerBase
    {
        int layer_id = 0;
        LayerType type;
        virtual ~LayerBase() = default;
        virtual void run(class ModelBuilder &builder, Pointcloud_utils &ptu) = 0;
        virtual void run(class ModelBuilder &builder, Pointcloud_online_utils &ptu) = 0;
        virtual void UpdateInput(ModelBuilder &builder, Pointcloud_utils &ptu) = 0;
        virtual void UpdateInput(ModelBuilder &builder, Pointcloud_online_utils &ptu) = 0;
        // virtual void ExportResult(ModelBuilder &builder) = 0;
    };

    /*
    struct LayerConfig
    {
        int layer_id = 0;
        uint8_t cfg_inch_loop_num = 3;
        uint32_t cfg_start_pt_addr = 0;
        uint8_t cfg_shift_num = 0;
        uint8_t cfg_pe_relu_enable = 0;
        uint8_t cfg_cam_search_range[3] = {1, 1, 1};
        uint8_t cfg_cnn_size[3] = {3, 3, 3};
        uint8_t cfg_bias_enable = 0;
        uint8_t cfg_cam_enable = 0x7;
        uint8_t cfg_core_enable = 0xf;

        uint8_t bias_core0[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        uint8_t bias_core1[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        uint8_t bias_core2[8] = {0, 0, 0, 0, 0, 0, 0, 0};
        uint8_t bias_core3[8] = {0, 0, 0, 0, 0, 0, 0, 0};


        uint64_t BASE_WT_ADDR;
        uint16_t DMA_WT_NUM = 4096; // 8*8*4*4*32*8/64 = 4096
        std::vector<uint64_t> wt_data;
    };
    */
    struct ModelBuilder;

    struct SubmConvLayerConfig : public LayerBase
    {
        LayerConfig layer_config;
        uint8_t cfg_och_loop_num; 

        uint8_t relu_enable = 1;
        uint64_t bias_cores[4] = {0, 0, 0, 0};          // bias for each core

        SubmConvLayerConfig(const LayerConfig &config) : layer_config(config)
        {
            type = LayerType::SUBMCONV;
            cfg_och_loop_num = get_used_core_num(config.cfg_core_enable) - 1;

            auto pack_signed_bias = [](const int8_t bias_arr[8]) -> uint64_t {
                uint64_t val = 0;
                for (int i = 0; i < 8; ++i) {
                    val |= (uint64_t(uint8_t(bias_arr[i])) << (i * 8));
                }
                return val;
            };

            bias_cores[0] = pack_signed_bias(layer_config.bias_core0);
            bias_cores[1] = pack_signed_bias(layer_config.bias_core1);
            bias_cores[2] = pack_signed_bias(layer_config.bias_core2);
            bias_cores[3] = pack_signed_bias(layer_config.bias_core3);

            relu_enable = layer_config.cfg_pe_relu_enable;
            layer_config.cfg_pe_relu_enable = 0;         // Disable on-chip ReLU
        }

        void run(ModelBuilder &builder, Pointcloud_utils &ptu) override;
        void run(ModelBuilder &builder, Pointcloud_online_utils &ptu) override;
        void UpdateInput(ModelBuilder &builder, Pointcloud_utils &ptu) override;
        void UpdateInput(ModelBuilder &builder, Pointcloud_online_utils &ptu) override;
    };

    struct FCLayerConfig : public LayerBase
    {
        std::vector<float> weights;
        std::vector<float> bias;
        uint16_t FC_ICH_LOOP_NUM = 0;
        uint16_t FC_OCH_LOOP_NUM = 0;

        FCLayerConfig() { type = LayerType::FC; }
        void run(ModelBuilder &builder, Pointcloud_utils &ptu) override;
        void run(ModelBuilder &builder, Pointcloud_online_utils &ptu) override;
        void UpdateInput(ModelBuilder &builder, Pointcloud_utils &ptu) override;
        void UpdateInput(ModelBuilder &builder, Pointcloud_online_utils &ptu) override;
        void fc_gemm(ModelBuilder &builder);
    };
    
    struct ModelBuilder
    {
        bool offline_mode = true;                    // simulate on-chip computing

        std::vector<LayerBase *> layers;

        int input_pt_num = 0;
        BlockConfig input;
        std::vector<BlockConfig> bcfg_partition;        
        std::vector<std::vector<int>> pt_mapping;
        std::vector<std::vector<int>> valid_filter;

        BlockConfig dummy_bcfg;                      // dummy block config for instruction initialization
        std::vector<uint64_t> merged_res;
        std::vector<int> res;

        uint64_t DDR_BASE_ADDR;
        uint64_t BASE_PT_ADDR;
        uint64_t BASE_IN_ADDR;
        uint64_t BASE_RS_ADDR;

        ModelBuilder(uint64_t ddr_addr,
                     uint64_t pt_addr,
                     uint64_t in_addr,
                     uint64_t rs_addr,
                     bool offline = true,
                     bool bias = false)
            : offline_mode(offline),
              DDR_BASE_ADDR(ddr_addr),
              BASE_PT_ADDR(pt_addr),
              BASE_IN_ADDR(in_addr),
              BASE_RS_ADDR(rs_addr)
        {
            MemInit();
        }

        ~ModelBuilder()
        {
            for (auto layer : layers)
            {
                delete layer;
            }
            layers.clear();

            if (input.spatial_shape)
            {
                delete[] input.spatial_shape;
                input.spatial_shape = nullptr;
            }
        }

        void addSubmConvLayer(const LayerConfig &config)
        {
            layers.push_back(new SubmConvLayerConfig(config));
        }

        void addFCLayer(const std::vector<float> &weights, const std::vector<float> &bias, uint32_t ICH_LOOP_NUM, uint32_t OCH_LOOP_NUM)
        {
            FCLayerConfig *fc_layer = new FCLayerConfig();
            fc_layer->weights = weights;
            fc_layer->bias = bias;
            fc_layer->FC_ICH_LOOP_NUM = ICH_LOOP_NUM;
            fc_layer->FC_OCH_LOOP_NUM = OCH_LOOP_NUM;
            layers.push_back(fc_layer);
        }

        // Initialize memory configurations
        void MemInit()
        {
            dummy_bcfg.cfg_cam_valid_num = 1024;
            dummy_bcfg.cfg_total_point_num = 1024;
            dummy_bcfg.BASE_PT_ADDR = BASE_PT_ADDR;
            dummy_bcfg.BASE_IN_ADDR = BASE_IN_ADDR;
            dummy_bcfg.BASE_RS_ADDR = BASE_RS_ADDR;

            input.BASE_PT_ADDR = BASE_PT_ADDR;
            input.BASE_IN_ADDR = BASE_IN_ADDR;
            input.BASE_RS_ADDR = BASE_RS_ADDR;
            input.spatial_shape = new uint16_t[6]{0, 0, 0, 0, 0, 0};
        }

        // Load model configuration from a JSON file
        void LoadModelConfig(const std::string &model_config_path);

        // Preprocess the input point cloud data
        void PreprocessInput();

        // Run model
        void RunModel();

        // Run online model
        void RunOnlineModel();

        // Run offline model
        void RunOfflineModel();
    };
}